import os
from apps.api.model.model.LSTM import LSTM
import torch
import torch.optim.lr_scheduler as lr_scheduler
from ruamel.yaml import YAML
from apps.api.model.dataset import CustomScaler, get_dataset, get_dataloader
import numpy as np
from tqdm import tqdm
import math
from sklearn.metrics import mean_squared_error, mean_absolute_percentage_error
import plotly.graph_objects as go


class main():
    def __init__(self, config, mode=None):
        # args为参数文件路径列表
        self.yaml = YAML(typ='safe', pure=True)
        with open(config, 'r', encoding='utf-8') as file:
            self.arg = self.yaml.load(file)
        self.mode = mode if mode is not None else self.arg['mode']
        self.epochs = None
        self.model = None
        self.criterion = None
        self.optim = None
        self.scheduler = None
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        # 实例化模型
        self.model = self.get_model(self.arg)
        # 实例化损失函数、优化器、学习器
        # if self.mode == 'train':
        self.epochs = self.arg['epochs']
        self.criterion, self.optim, self.scheduler = self.get_train_instance(self.arg)
        # 获取dataset、dataloader
        self.trainset, self.valset = get_dataset(self.arg)
        self.trainloader, self.valloader = get_dataloader(self.trainset, self.valset, self.arg)
        # 获取scaler
        self.scaler = self.trainset.return_scaler()
        # 获取数据集名字
        self.dataset_name = '-'.join(self.arg['pred_col'])
        # 数据集总长度(包含训练集、测试集)
        self.data_len = self.trainset.datas_len  # 返回源数据总长度
        # 价格列索引及列名
        self.price_col_index = self.trainset.get_price_col_index()
        self.price_col = self.arg['price_col']
        # 获取更新scaler参数後的arg
        self.arg = self.trainset.return_arg()
        # 其他参数
        self.save_fre = self.arg['save_fre']

    def get_model(self, arg):
        model = LSTM(input_dim=arg['input_dim'], hidden_dim=arg['hidden_dim'], output_dim=arg['output_dim'],
                     num_layers=arg['num_layers']).to(self.device)
        if self.mode == 'train' and arg['pre_ckpt']:
            model.load_state_dict(torch.load(arg['pre_ckpt'], map_location=self.device))
        elif self.mode == 'prediction':
            assert arg['ckpt'] is not None, '未输入模型'
            model.load_state_dict(torch.load(arg['ckpt'], map_location=self.device))
        return model

    def get_train_instance(self, arg):
        # 损失函数
        criterion = torch.nn.MSELoss().to(self.device)
        # 优化器
        if arg['optim'] == 'Adam':
            optimizer = torch.optim.Adam(self.model.parameters(), lr=float(arg['lr']))
        elif arg['optim'] == 'SGD':
            optimizer = torch.optim.SGD(self.model.parameters(), lr=float(arg['lr']), momentum=arg['momentum'],
                                        weight_decay=arg['weight_decay'])
        else:
            raise ValueError("请检查配置文件中optim是否正确，你的输入为{}，应为Adam or SGD".format(arg['optim']))
        # scheduler
        if arg['scheduler']:
            scheduler = lr_scheduler.StepLR(optimizer, step_size=arg['step_size'], gamma=arg['gamma'])
        else:
            scheduler = None
        return (criterion, optimizer, scheduler)

    def train(self):
        self.savedir = self.creat_savedir(self.arg['savedir'])
        train_loss = np.zeros(self.epochs)
        val_loss = np.zeros(self.epochs)
        best = (1, 0, 0)
        for epoch in tqdm(range(self.epochs)):
            # 训练
            totle_loss_train, train_pred, train_label = self.train_step(self.trainloader)
            train_pred, train_label, train_metrics = self.get_metrics(train_pred, train_label)
            train_loss[epoch] = totle_loss_train
            # 验证
            totle_loss_val, val_pred, val_label = self.val_step(self.valloader)
            val_loss[epoch] = totle_loss_val
            val_pred, val_label, val_metrics = self.get_metrics(val_pred, val_label)

            # 输出指标
            print(f"\ntrain_step:MSE{totle_loss_train}------")
            for k, v in train_metrics.items():
                print(f'{k}:  RMSE:{v["RMSE"]:.3f},MAPE:{v["MAPE"]:.3f}')
            print(f"\nval_step:  MSE{totle_loss_train}------")
            for k, v in val_metrics.items():
                print(f'{k}:  RMSE:{v["RMSE"]:.3f},MAPE:{v["MAPE"]:.3f}')

            # 权重保存规则
            MAPE = val_metrics['Avg']['MAPE']
            if epoch % self.save_fre == 0:
                name = f'epoch{epoch}-loss{totle_loss_val:.2f}-MAPE{MAPE:.2f}.pth'
                torch.save(self.model.state_dict(), os.path.join(self.savedir, name))
            if MAPE < best[0]:
                # 获取best权重
                best_name = f'best-epoch{epoch}-loss{totle_loss_val:.2f}-MAPE{MAPE * 100:.2f}%.pth'
                best = (MAPE, self.model.state_dict(), best_name)
            if epoch == self.epochs - 1:
                # 保存last权重
                name = f'last-epoch{epoch}-loss{totle_loss_val:.2f}-MAPE{MAPE:.2f}.pth'
                torch.save(self.model.state_dict(), os.path.join(self.savedir, name))
                # 保存best权重
                print(f'best_MAPE:{best[0]}')
                torch.save(best[1], os.path.join(self.savedir, best[2]))
                # 保存yaml文件
                # ruamel.yaml.scalarstring.walk_tree(str)
                # ruamel.yaml.scalarstring.DEFAULT_ENCODING = 'utf-8'
                with open(os.path.join(self.savedir, f"config.yaml"), 'w', encoding='utf-8') as file:
                    self.yaml.dump(self.arg, file)

        plot_data, residuals = self.creat_plot_data(train_label, val_label, train_pred, val_pred)
        for i in range(plot_data.shape[2]):
            self.plot(plot_data[:, :, i], self.price_col[i])

    def predict(self, show=False):
        # 训练
        _, train_pred, train_label = self.val_step(self.trainloader)
        train_pred, train_label, train_metrics = self.get_metrics(train_pred, train_label)
        # 验证
        _, val_pred, val_label = self.val_step(self.valloader)
        val_pred, val_label, val_metrics = self.get_metrics(val_pred, val_label)
        # 输出指标
        for k, v in train_metrics.items():
            print(f'\n{k}:  RMSE:{v["RMSE"]:.3f},MAPE:{v["MAPE"]:.3f}')
        for k, v in val_metrics.items():
            print(f'{k}:  RMSE:{v["RMSE"]:.3f},MAPE:{v["MAPE"]:.3f}')
        # 获取绘图数据
        plot_data, residuals = self.creat_plot_data(train_label, val_label, train_pred, val_pred)
        if show:
            # 绘图
            for i in range(plot_data.shape[2]):
                self.plot(plot_data[:, :, i], self.price_col[i])
        return plot_data, self.trainset.row_datas['pub_time'], np.round(residuals[:, 0], 2).tolist()

    def creat_plot_data(self, train_label, val_label, train_pred, val_pred):
        lable = np.vstack((train_label, val_label))  # (data_len,price_len)
        pred = np.vstack((train_pred, val_pred))
        residuals = lable - pred
        datas_len = self.data_len  # 总数据长度
        lable_len = lable.shape[0]  # 真实值(标签数据长度)
        price_len = len(self.price_col)  # 预测的价格列
        seg_len = self.arg['seg_len']
        plot_data = np.full((3, datas_len, price_len),
                            np.nan)  # shape = ((label,train_pred,val_pred),data_len,price_len)
        plot_data[0, seg_len:, :] = lable
        plot_data[1, seg_len:len(train_pred) + seg_len, :] = train_pred
        plot_data[2, len(train_pred) + seg_len:len(train_pred) + seg_len + len(val_pred), :] = val_pred
        return plot_data, residuals

    def plot(self, plot_data, title):
        fig = go.Figure()
        x = np.arange(plot_data.shape[1])
        fig.add_trace(go.Scatter(go.Scatter(x=x, y=plot_data[0],
                                            mode='lines',
                                            name='Actual Value')))
        fig.add_trace(go.Scatter(go.Scatter(x=x, y=plot_data[1],
                                            mode='lines',
                                            name='Train prediction')))
        fig.add_trace(go.Scatter(x=x, y=plot_data[2],
                                 mode='lines',
                                 name='Test prediction'))

        fig.update_layout(
            xaxis=dict(
                showline=True,
                showgrid=True,
                showticklabels=False,
                linecolor='white',
                linewidth=2
            ),
            yaxis=dict(
                title_text='Price(RMB)',
                titlefont=dict(
                    family='Rockwell',
                    size=12,
                    color='white',
                ),
                showline=True,
                showgrid=True,
                showticklabels=True,
                linecolor='white',
                linewidth=2,
                ticks='outside',
                tickfont=dict(
                    family='Rockwell',
                    size=12,
                    color='white',
                ),
            ),
            showlegend=True,
            template='plotly_dark'

        )

        annotations = []
        annotations.append(dict(xref='paper', yref='paper', x=0.0, y=1.05,
                                xanchor='left', yanchor='bottom',
                                text=f'Results ({title})',
                                font=dict(family='Rockwell',
                                          size=26,
                                          color='white'),
                                showarrow=False))
        fig.update_layout(annotations=annotations)
        fig.show()

    def creat_savedir(self, savedir):
        assert os.path.isdir(savedir)
        os.makedirs(savedir, exist_ok=True)
        li = [int(i) for i in os.listdir(savedir)]
        li.sort()
        step = int(li[-1]) + 1 if li else 0
        savedir = savedir + '/' + str(step)
        os.makedirs(savedir, exist_ok=True)
        return savedir

    def train_step(self, loader):
        self.model.train()
        totle_loss = 0
        train_pred = []
        train_label = []
        for datas, labels in loader:
            pred = self.model(datas)  # (277,1)
            loss = self.criterion(pred, labels)

            self.optim.zero_grad()
            loss.backward()
            self.optim.step()
            #
            totle_loss += loss.item()
            train_pred.append(pred.detach().numpy())
            train_label.append(labels.detach().numpy())
        if self.scheduler is not None:
            self.scheduler.step()
        return totle_loss, train_pred, train_label

    def val_step(self, loader):
        self.model.eval()
        totle_loss = 0
        val_pred = []
        val_label = []
        for datas, labels in loader:
            with torch.no_grad():
                pred = self.model(datas)
                loss = self.criterion(pred, labels)
                totle_loss += loss.item()
                val_pred.append(pred.detach().numpy())
                val_label.append(labels.detach().numpy())
        return totle_loss, val_pred, val_label

    def get_metrics(self, preds, labels):
        def metrics_step(preds, labels):
            RMSE = math.sqrt(mean_squared_error(labels, preds))
            # 计算MAPE
            MAPE = mean_absolute_percentage_error(labels, preds)
            return RMSE, MAPE

        preds = self.scaler.inverse_transform(np.concatenate(preds))[:, self.price_col_index]
        labels = self.scaler.inverse_transform(np.concatenate(labels))[:, self.price_col_index]
        metrics = {}
        totle_MAPE = 0
        totle_RMSE = 0
        num = preds.shape[1]
        for i in range(num):
            RMSE, MAPE = metrics_step(preds[:, i], labels[:, i])
            metrics[self.price_col[i]] = {'RMSE': RMSE, 'MAPE': MAPE}
            totle_MAPE += MAPE
            totle_RMSE += RMSE
        metrics['Avg'] = {'RMSE': totle_RMSE / num, 'MAPE': totle_MAPE / num}
        return preds, labels, metrics


if __name__ == '__main__':
    # config = r'./config/default.yaml'
    # model = main(config,'train')
    # model.train()
    #
    config = r'./checkpoint/1/config.yaml'
    model = main(config, 'prediction')
    model.predict()
